import math
import open3d as o3d
import open3d as o3d
import numpy as np
import copy
import sys
import cv2
np.set_printoptions(suppress=True)


def gen_cloud():

    cv_file = cv2.FileStorage(
        "../camera.yml", cv2.FILE_STORAGE_READ)  # 实例化一个 FileStorage
    # getNode 成员函数将获得一个 FileNode 类型的对象，mat() 函数将 FileNode 转换为矩阵
    # trans_init = cv_file.getNode("CameraExtrinsicMat").mat()
    camk = cv_file.getNode("camera_matrix").mat()
    # center = np.array([trans_init[0, 3], trans_init[1, 3], trans_init[2, 3]])

    # trans_init = np.linalg.inv(trans_init)
    # trans_init = np.loadtxt("rt")
    # print(trans_init)
    print(camk, "\n")

    intrinsic = o3d.camera.PinholeCameraIntrinsic(
        640, 480, camk[0, 0], camk[0, 0],  camk[0, 2], camk[1, 2])
    print(intrinsic)
    # print(camk[0, 0], camk[0, 0],  camk[0, 2], camk[1, 0])
    filepath = sys.argv[1]

    pose = cv2.FileStorage(f'{filepath}/pose.yml', cv2.FileStorage_READ)
    print(f'{filepath}/pose.yml')
    axis_pcd = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0, origin=[
        0, 0, 0])
    idsx = pose.getNode('idx').mat()
    # print(idsx)
    pcds = [axis_pcd]
    for [i] in idsx:
        #     if i > 50:
        #         break
        # print(i)
        color_raw = o3d.io.read_image(f"{filepath}/rgb/{i}.png")
        depth_raw = o3d.io.read_image(f"{filepath}/depth/{i}d.png")
        # print(color_raw,depth_raw)
        rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(
            color_raw, depth_raw, depth_scale=1000, depth_trunc=50.0, convert_rgb_to_intensity=False)
        transformation = pose.getNode('p'+str(i)).mat()
        print(transformation)
        pcd = o3d.geometry.PointCloud.create_from_rgbd_image(
            rgbd_image, intrinsic)
    #     o3d.io.write_point_cloud("test.ply",pcd,write_ascii=True)
    #     break
    #     pcd = pcd.voxel_down_sample(0.5)
    #     # Flip it, otherwise the pointcloud will be upside down
        # pcd.transform([ [0, 1, 0, 0], [1, 0, 0, 0],[0, 0, -1, 0], [0, 0, 0, 1]])
        pcd.transform(transformation)
        pcds.append(pcd)
    ps = pcds[1]
    for i, c in enumerate(pcds):
        if i < 2:
            continue
        ps += c
    ps = ps.voxel_down_sample(0.001)
    o3d.io.write_point_cloud("test.ply", ps)
    print("gen test.ply")
    # o3d.visualization.draw_geometries(pcds)


def scalable_integrate_rgb_frames(path_dataset, intrinsic):

    pas = np.loadtxt(path_dataset+"/file.txt", np.str_)
    # print(pas)
    pose = cv2.FileStorage(f'{path_dataset}/pose.yml', cv2.FileStorage_READ)
    print(f'{path_dataset}/pose.yml')
    axis_pcd = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0, origin=[
        0, 0, 0])
    idsx = pose.getNode('idx').mat()

    volume = o3d.integration.ScalableTSDFVolume(
        voxel_length=0.002,
        sdf_trunc=0.03,
        color_type=o3d.integration.TSDFVolumeColorType.RGB8)

    for idx, [i] in enumerate(idsx):
        # print(path_dataset+pas[idx][1],i,idx)
        color_raw = o3d.io.read_image(path_dataset+pas[i][1])
        depth_raw = o3d.io.read_image(path_dataset+pas[i][0])
        # print(f"{path_dataset}/rgb/{i}.png",color_raw,depth_raw)
        rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(
            color_raw, depth_raw, depth_scale=1000, depth_trunc=50.0, convert_rgb_to_intensity=False)
        transformation = pose.getNode('p'+str(i)).mat()
        # print(transformation)
        volume.integrate(rgbd_image, intrinsic, np.linalg.inv(transformation))

    mesh = volume.extract_triangle_mesh()
    mesh.compute_vertex_normals()
    o3d.io.write_triangle_mesh(path_dataset+"/mesh.ply", mesh, False, True)
    # o3d.visualization.draw_geometries([mesh])


cv_file = cv2.FileStorage(
    "../camera.yml", cv2.FILE_STORAGE_READ)  # 实例化一个 FileStorage
filepath = sys.argv[1]
camk = cv_file.getNode("camera_matrix").mat()
print(camk, "\n")
intrinsic = o3d.camera.PinholeCameraIntrinsic(
    640, 480, camk[0, 0], camk[0, 0],  camk[0, 2], camk[1, 2])
scalable_integrate_rgb_frames(filepath, intrinsic)
